import cv2
import cv2 as cv
from ultralytics import YOLO
import math
import time
import threading
import queue
import datetime

# 加载模型
pose_model = YOLO('yolov8n-pose.pt')  # 姿态估计模型
phone_model = YOLO('yolov8n.pt')  # 手机检测模型
smoking_model = YOLO('bests.pt')  # 吸烟检测模型

# 创建一个队列来存储帧
frame_queue = queue.Queue()
stop_event = threading.Event()


# 读取帧的线程
def read_frames(cap, frame_queue, stop_event):
    while not stop_event.is_set():
        ret, frame = cap.read()
        if not ret:
            stop_event.set()
            break
        frame_queue.put(frame)

# 处理帧的线程
def process_frames(frame_queue, stop_event):
    epoch = 0
    while not stop_event.is_set() or not frame_queue.empty():
        try:
            frame = frame_queue.get(timeout=1)  # 设置超时，以防队列为空时阻塞
        except queue.Empty:
            continue

        # 检测抽烟
        smoking_detected = False
        smoking_result = smoking_model(frame)  # 使用模型推理
        smoking_result = smoking_result[0]
        smoking_boxes = smoking_result.boxes
        smoking_class_ids = smoking_boxes.cls.cpu().numpy().astype(int)

        # 确保类别 0 是抽烟行为
        for smoking_box in smoking_boxes:
            class_id = int(smoking_box.cls.cpu().numpy())
            if class_id == 0:  # 假设类别 0 表示抽烟
                smoking_detected = True
                # 获取边界框坐标
                x1, y1, x2, y2 = map(int, smoking_box.xyxy[0].cpu().numpy())
                # 在图片上绘制边界框
                cv.rectangle(frame, (x1, y1), (x2, y2), (0, 0, 255), 2)  # 红色框，线宽为2
                cv.putText(frame, "Smoking", (x1, y1 - 10), cv.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2)  # 添加标签
                break

        # 检测手机
        phone_detected = False
        phone_result = phone_model(frame)
        phone_result = phone_result[0]
        phone_boxes = phone_result.boxes
        phone_class_ids = phone_boxes.cls.cpu().numpy().astype(int)

        # 确保类别 67 是手机
        if 67 in phone_class_ids:  # 检测到手机
            for phone_box in phone_boxes:
                xyxy_tensor = phone_box.xyxy[0]  # 提取cellphone边界框的xyxy信息
                xyxy_coords = xyxy_tensor.cpu().numpy()
                x1, y1, x2, y2 = xyxy_coords
                phone_detected = True

        # 检测姿态（仅在检测到手机时进行）
        if phone_detected:
            pose_results = pose_model(frame)
            keypoints = pose_results[0].keypoints

            xy_cpu = keypoints.xy.cpu() if keypoints.xy.device != 'cpu' else keypoints.xy

            # 提取点位
            shoulder_x_r, shoulder_y_r = xy_cpu[0, 6, 0].item(), xy_cpu[0, 6, 1].item()
            elbow_x_r, elbow_y_r = xy_cpu[0, 8, 0].item(), xy_cpu[0, 8, 1].item()
            wrist_x_r, wrist_y_r = xy_cpu[0, 10, 0].item(), xy_cpu[0, 10, 1].item()
            shoulder_x_l, shoulder_y_l = xy_cpu[0, 5, 0].item(), xy_cpu[0, 5, 1].item()
            elbow_x_l, elbow_y_l = xy_cpu[0, 7, 0].item(), xy_cpu[0, 7, 1].item()
            wrist_x_l, wrist_y_l = xy_cpu[0, 9, 0].item(), xy_cpu[0, 9, 1].item()

            hip_y_l = xy_cpu[0, 11, 1]
            hip_y_r = xy_cpu[0, 12, 1]

            # 计算向量和角度等
            vector_A_x_r = elbow_x_r - shoulder_x_r
            vector_A_y_r = elbow_y_r - shoulder_y_r
            vector_B_x_r = wrist_x_r - elbow_x_r
            vector_B_y_r = wrist_y_r - elbow_y_r
            vector_A_x_l = elbow_x_l - shoulder_x_l
            vector_A_y_l = elbow_y_l - shoulder_y_l
            vector_B_x_l = wrist_x_l - elbow_x_l
            vector_B_y_l = wrist_y_l - elbow_y_l

            dot_product_r = vector_A_x_r * vector_B_x_r + vector_A_y_r * vector_B_y_r
            dot_product_l = vector_A_x_l * vector_B_x_l + vector_A_y_l * vector_B_y_l

            magnitude_A_r = math.sqrt(vector_A_x_r ** 2 + vector_A_y_r ** 2)
            magnitude_B_r = math.sqrt(vector_B_x_r ** 2 + vector_B_y_r ** 2)
            magnitude_A_l = math.sqrt(vector_A_x_l ** 2 + vector_A_y_l ** 2)
            magnitude_B_l = math.sqrt(vector_B_x_l ** 2 + vector_B_y_l ** 2)

            maxmagnitude = max(magnitude_B_l, magnitude_B_r, magnitude_A_l, magnitude_A_r)

            if magnitude_A_r * magnitude_B_r == 0:
                continue
            if magnitude_A_l * magnitude_B_l == 0:
                continue

            cosine_angle_r = dot_product_r / (magnitude_A_r * magnitude_B_r)
            cosine_angle_l = dot_product_l / (magnitude_A_l * magnitude_B_l)

            closest_angle = min(abs(cosine_angle_l), abs(cosine_angle_r))

            x_phone_mid = (x1 + x2) / 2
            y_phone_mid = (y1 + y2) / 2

            phone_to_leftwrist = math.sqrt((x_phone_mid - wrist_x_l) ** 2 + (y_phone_mid - wrist_y_l) ** 2)
            phone_to_rightwrist = math.sqrt((x_phone_mid - wrist_x_r) ** 2 + (y_phone_mid - wrist_y_r) ** 2)
            phone_to_wirst = min(phone_to_rightwrist, phone_to_leftwrist)

            height_l = shoulder_y_l - hip_y_l
            height_r = shoulder_y_r - hip_y_r
            delta_y_hiptoshoulder = max(height_l, height_r)
            delta_x_shoulder = math.sqrt((shoulder_x_r - shoulder_x_l) ** 2 + (shoulder_y_r - shoulder_y_l) ** 2)
            bili = delta_x_shoulder / delta_y_hiptoshoulder

            # 检查条件
            if -0.5345 - 10 <= closest_angle <= 0.2756 + 10 and phone_to_wirst <= maxmagnitude:
                if abs(bili) <= 0.36149454 + 0.1353443 * 5.5:
                    phone_detected = True

        # 保存违规行为图片（统一命名）
        if smoking_detected or phone_detected:
            cv.imwrite(f'F:/电子科大/&【大三上】/挑战性课程/静态/images/violation_{epoch}.jpg', frame)
            epoch += 1
            print("检测到违规行为，图片已保存")

        # 显示图像
        cv.imshow("image", frame)
        if cv.waitKey(1) == ord("q"):
            stop_event.set()

if __name__ == "__main__":
    # 使用默认电脑摄像头.
    cap = cv2.VideoCapture("http://192.168.95.250:4747/mjpegfeed")
    #("http://192.168.20.250:4747/mjpegfeed")
    #"http://192.168.81.21:4747/mjpegfeed"
    #cap = cv.VideoCapture("http://192.168.217.250:4747/mjpegfeed")

    # 使用外置摄像头
    # cap = cv2.VideoCapture("http://192.168.217.250:4747/mjpegfeed")

    # 启动读取帧的线程
    reader_thread = threading.Thread(target=read_frames, args=(cap, frame_queue, stop_event))
    reader_thread.start()

    # 启动处理帧的线程
    processor_thread = threading.Thread(target=process_frames, args=(frame_queue, stop_event))
    processor_thread.start()

    # 等待线程完成
    reader_thread.join()
    processor_thread.join()

    # 释放资源
    cap.release()
    cv.destroyAllWindows()